package edu.cornell.cs4740.qa.parsing;

import java.io.BufferedReader;
import java.io.DataInputStream;
import java.io.FileInputStream;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.List;

import edu.cornell.cs4740.qa.Document;
import edu.cornell.cs4740.qa.Question;
import edu.cornell.cs4740.qa.Sentence;
import edu.cornell.cs4740.qa.WordContainer;
import edu.cornell.cs4740.qa.baseline.Answer;

public class ParseDocsFile {
  private static enum State {
    NONE, DOC, CAPTION, LEADPARA, HEADLINE, TEXT
  }
  
  public static List<Document> parseDocsFiles(String filepath) throws Exception {
    String[] temp = filepath.split("/");
    String file = temp[temp.length - 1];
    FileInputStream fstream = new FileInputStream(filepath);
    DataInputStream in = new DataInputStream(fstream);
    BufferedReader br = new BufferedReader(new InputStreamReader(in));
    String line;
    
    List<Document> documents = new ArrayList<Document>();
    State curState = State.NONE;
    StringBuilder textBuilder = new StringBuilder();
    String qid = "";
    int rank = 0;
    double score = 0;
    
    while ((line = br.readLine()) != null) {
      switch (curState) {
      
      case NONE:
        if (line.startsWith("Qid")) {
          String[] lineInfo = line.split("\t");
          qid = lineInfo[0].split(": ")[1];
          rank = Integer.parseInt(lineInfo[1].split(": ")[1]);
          score = Double.parseDouble(lineInfo[2].split(": ")[1]);
          curState = State.DOC;
        }
        break;
        
      case DOC:
        if (line.contains("</DOC>")) {
          Document doc = new Document(file, qid, rank, score);
          doc.addDocText(textBuilder.toString());
          
          TokenFactory tokenFactory = new TokenFactory();
          List<String> sentenceStrings = tokenFactory.getSentences(textBuilder.toString());
          for(String st : sentenceStrings) {
            List<WordContainer> tokenList = tokenFactory.sentenceToTokens(st);
            doc.addSentence(new Sentence(tokenList));
            doc.addSentenceString(st);
          }
          documents.add(doc);

          // reset everything
          curState = State.NONE;
          textBuilder = new StringBuilder();
          qid = "";
          rank = 0;
          score = 0;
        } else if (line.startsWith("<")) {
          // look for info start tags
          if (line.startsWith("<CAPTION>")) {
            if (line.endsWith("</CAPTION>") && !line.equals("")) {
              textBuilder.append(line.substring(9, line.length() - 10));
              break;
            }
            textBuilder.append(line.substring(9));
            curState = State.CAPTION;
          } else if (line.startsWith("<LEADPARA>")) {
            if (line.endsWith("</LEADPARA>") && !line.equals("")) {
              textBuilder.append(line.substring(10, line.length() - 11));
              break;
            }
            textBuilder.append(line.substring(10));
            curState = State.LEADPARA;
          } else if (line.startsWith("<HEADLINE>")) {
            if (line.endsWith("</HEADLINE>") && !line.equals("")) {
              textBuilder.append(line.substring(10, line.length() - 11));
              break;
            }
            textBuilder.append(line.substring(10));
            curState = State.HEADLINE;
          } else if (line.startsWith("<TEXT>")) {
            if (line.endsWith("</TEXT>") && !line.equals("")) {
              textBuilder.append(line.substring(6, line.length() - 7));
              break;
            }
            textBuilder.append(line.substring(6));
            curState = State.TEXT;
          }
        }
        break;
        
      case CAPTION:
        if (line.endsWith("</CAPTION>") && !line.equals("")) {
          textBuilder.append(line.substring(0, line.length() - 10));
          curState = State.DOC;
        } else {
          textBuilder.append(line);
        }
        break;
        
      case LEADPARA:
        if (line.endsWith("</LEADPARA>") && !line.equals("")) {
          textBuilder.append(line.substring(0, line.length() - 11));
          curState = State.DOC;
        } else {
          textBuilder.append(line);
        }
        break;
        
      case HEADLINE:
        if (line.endsWith("</HEADLINE>") && !line.equals("")) {
          textBuilder.append(line.substring(0, line.length() - 11));
          curState = State.DOC;
        } else {
          textBuilder.append(line);
        }
        break;
        
      case TEXT:
        if (line.endsWith("</TEXT>") && !line.equals("")) {
          textBuilder.append(line.substring(0, line.length() - 7));
          curState = State.DOC;
        } else {
          textBuilder.append(line);
        }
        break;
      }
    }
    
    in.close();
    
    return documents;
  }
  
  public static List<Question> parseQuestionFile(String filepath) throws Exception {
    FileInputStream fstream = new FileInputStream(filepath);
    DataInputStream in = new DataInputStream(fstream);
    BufferedReader br = new BufferedReader(new InputStreamReader(in));
    String line;
    String questionNum = null;
    boolean descNext = false;
    
    List<Question> questions = new ArrayList<Question>();
    
    while ((line = br.readLine()) != null) {
      if (line.contains("<num>")) {
        String[] idLine = line.split(" ");
        questionNum = idLine[2];
      } else if (line.contains("<desc>")) {
        descNext = true;
      } else if (descNext) {
        descNext = false;
        Question question = new Question(questionNum, line);
        questions.add(question);
        questionNum = null;
      }
    }
    
    in.close();
    
    return questions;
  }
  
  public static List<Answer> parseAnswerFile(String filepath) throws Exception {
    FileInputStream fstream = new FileInputStream(filepath);
    DataInputStream in = new DataInputStream(fstream);
    BufferedReader br = new BufferedReader(new InputStreamReader(in));
    String line;
    
    String questionNum = null;
    String answer = null;
    int lineNum = 0; // 0 - Question id, 3 - Answer
    
    List<Answer> answers = new ArrayList<Answer>();
    
    while ((line = br.readLine()) != null) {
      if (lineNum == 0) {
        if (line.toLowerCase().startsWith("question")) {
          String[] split = line.split(" ");
          questionNum = split[1];
          lineNum++;
        }
      } else if (lineNum < 3) {
        lineNum++;
      } else if (lineNum == 3) {
        answer = line;
        Answer a = new Answer(questionNum, null, answer);
        answers.add(a);
        lineNum = 0;
      }
    }
    
    in.close();
    
    return answers;
  }
}
